package turritopsis

import (
    "sync"
    "adkr/pkg/adkr"
)

// InMemoryRouter is a simple NetRouter implementation using Go channels.
// It is intended for single-process tests and local integration runs.
type InMemoryRouter struct {
    mu sync.Mutex
    // per-epoch channels for KR traffic
    krIn  map[int]chan adkr.Message
    krOut map[int]chan adkr.Message
    krFwd map[int]bool
    // optional: per-epoch channels for ACS (unused by default)
    acsIn  map[int]chan adkr.Message
    acsOut map[int]chan adkr.Message
    bufSize int
}

// NewInMemoryRouter creates a new in-memory router with the given channel buffer size.
func NewInMemoryRouter(bufferSize int) *InMemoryRouter {
    if bufferSize <= 0 { bufferSize = 128 }
    return &InMemoryRouter{
        krIn:   make(map[int]chan adkr.Message),
        krOut:  make(map[int]chan adkr.Message),
        krFwd:  make(map[int]bool),
        acsIn:  make(map[int]chan adkr.Message),
        acsOut: make(map[int]chan adkr.Message),
        bufSize: bufferSize,
    }
}

func (r *InMemoryRouter) KRChannels(epoch int) (<-chan adkr.Message, chan<- adkr.Message) {
    r.mu.Lock()
    defer r.mu.Unlock()
    in := r.krIn[epoch]
    out := r.krOut[epoch]
    if in == nil { in = make(chan adkr.Message, r.bufSize); r.krIn[epoch] = in }
    if out == nil { out = make(chan adkr.Message, r.bufSize); r.krOut[epoch] = out }
    if !r.krFwd[epoch] {
        r.krFwd[epoch] = true
        // forwarder: emulate loopback network
        go func(inch chan adkr.Message, outch chan adkr.Message) {
            for m := range outch {
                inch <- m
            }
        }(in, out)
    }
    return in, out
}

func (r *InMemoryRouter) ACSChannels(epoch int) (<-chan adkr.Message, chan<- adkr.Message) {
    r.mu.Lock()
    defer r.mu.Unlock()
    in := r.acsIn[epoch]
    out := r.acsOut[epoch]
    if in == nil { in = make(chan adkr.Message, r.bufSize); r.acsIn[epoch] = in }
    if out == nil { out = make(chan adkr.Message, r.bufSize); r.acsOut[epoch] = out }
    return in, out
}


